/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2019-2022. All rights reserved.
 * Description: 集合通信算子信息库头文件
 * Author: lilianlin
 * Create: 2019-11-28
 */

#ifndef AUTO_TUNING_HCOM_OPS_KERNEL_INFO_STORE_H
#define AUTO_TUNING_HCOM_OPS_KERNEL_INFO_STORE_H

#include "hcom_ops_kernel_info_store.h"
#include <map>
#include "mem_host_pub.h"

namespace hccl {
// Ge适配的类
class AutoTuningHcomOpsKernelInfoStore : public HcomOpsKernelInfoStore {
public:
    AutoTuningHcomOpsKernelInfoStore();
    ~AutoTuningHcomOpsKernelInfoStore() override;
    // load task for op
    ge::Status LoadTask(ge::GETaskInfo &task) override;
    ge::Status UnloadTask(ge::GETaskInfo &task) override;

private:
    HcclResult SetCustomKernelInfo(ge::OpInfo &opinfo, std::map<string, ge::OpInfo> &infos) const override;
    HcclResult CheckPrivateDef(const ge::GETaskInfo &task) override;
    HcclResult GetDataTypeFromTaskInfo(const ge::GETaskInfo &task, HcclDataType &dataType) const override;
    HcclResult GetRankSizeFromTaskInfo(const ge::GETaskInfo &task, uint32_t &rankSize);
    HcclResult GetHcclInfo(const ge::GETaskInfo &task, ge::GETaskKernelHcclInfo& hcclInfo);
    HcclResult HCCLOpsKernel(const ge::GETaskInfo &task, const std::string &sCollectiveType);
    HcclResult HcomBroadcastOpKernel(const ge::GETaskInfo &task);
    HcclResult HcomAllReduceOpKernel(const ge::GETaskInfo &task);
    HcclResult HcomAllGatherOpKernel(const ge::GETaskInfo &task);
    HcclResult HcomAlltoAllVOpKernel(const ge::GETaskInfo &task);
    HcclResult HcomAlltoAllVCOpKernel(const ge::GETaskInfo &task);
    HcclResult HcomAlltoAllOpKernel(const ge::GETaskInfo &task);
    HcclResult GetOriginalGraphShapeTypeFromTaskInfo(const ge::GETaskInfo &task, u32 &shapeType) override;
    // add
    HcclResult HcomReduceOpKernel(const ge::GETaskInfo &task);
    HcclResult HcomReduceScatterOpKernel(const ge::GETaskInfo &task);
    HcclResult HcomSendOpKernel(const ge::GETaskInfo &task);
    HcclResult HcomReceiveOpKernel(const ge::GETaskInfo &task);
    HcclResult InitAlltoAllHostMem(HcclDataType recvType, u64 recvCount, HostMem& hostMemEmpty);

    std::map<u64, std::map<HcclDataType, HostMem>> mapCountTypeHostMem_;
};
}
#endif // AUTO_TUNING_HCOM_OPS_KERNEL_INFO_STORE_H
